"""
@Description :   显示社区编号的网格图像，自动分配颜色
@Author      :   Li Junjie 
@Time        :   2025/03/20 10:36:47
"""

import cv2
import numpy as np
import json

def save_to_jsonl(data, file_path):
    with open(file_path, 'w', encoding='utf-8') as f:
        for item in data:
            json_line = json.dumps(item, ensure_ascii=False)
            f.write(json_line + '\n')

# 读取 JSON 文件
json_file = "data/node_community_colors.json"
with open(json_file, 'r', encoding='utf-8') as f:
    color_data = json.load(f)

# 解析颜色信息
color_map = {}
for label, info in color_data.items():
    hex_color = info["color"]  # 获取颜色（如 "#1f77b4"）
    # 将 HEX 颜色转换为 BGR 格式（OpenCV 颜色格式）
    bgr_color = tuple(int(hex_color[i:i+2], 16) for i in (5, 3, 1))  # 反转 RGB 顺序
    color_map[int(label)] = bgr_color

# 读取图像
image_path = 'tools/test.png'  # 替换为你的图像路径
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)  # 读取为灰度图

# 设置阈值并转换为二值图，确保黑色区域为255，白色区域为0
_, binary_image = cv2.threshold(image, 254, 255, cv2.THRESH_BINARY_INV)

# 查找黑色区域的最小外接矩形
contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
x, y, w, h = cv2.boundingRect(contours[0])  # 获取最小外接矩形坐标和尺寸

# 裁剪图像以仅保留黑色区域的最小外接矩形
cropped_image = binary_image[y:y+h, x:x+w]

# 定义每个分割块的大小
block_side = 50  # 设置每个网格的边长，例如 50x50 的网格

# 编号计数器
label_count = 1
segments = []  # 存储分割的矩形区域及其编号
positions = {}  # 用于存储编号和位置的映射关系

# 遍历裁剪后的图像，并按固定大小的网格划分
for j in range(0, cropped_image.shape[0], block_side):
    for i in range(0, cropped_image.shape[1], block_side):
        # 计算网格的坐标
        x1, y1 = i, j
        x2, y2 = min(i + block_side, cropped_image.shape[1] - 1), min(j + block_side, cropped_image.shape[0] - 1)
        
        # 检查是否大部分是黑色区域
        block = cropped_image[y1:y2, x1:x2]
        black_ratio = np.sum(block == 255) / (block.shape[0] * block.shape[1])
        
        if (y1-y2) == (x1-x2):
            if black_ratio > 0.7:  # 如果超过70%是黑色，则编号
                segments.append((label_count, (x1, y1, x2, y2)))
                # 存储编号和位置的映射
                positions[label_count] = (x1, y1)
                label_count += 1

# 创建结果图像以显示编号
output_image = cv2.cvtColor(cropped_image, cv2.COLOR_GRAY2BGR)
# 遍历每个分割的网格并填充颜色
for label, (x1, y1, x2, y2) in segments:
    color = color_map.get(label, (200, 200, 200))  # 获取颜色，默认灰色

    # 填充矩形区域
    cv2.rectangle(output_image, (x1, y1), (x2, y2), color, thickness=-1)
    cv2.rectangle(output_image, (x1, y1), (x2, y2), (1, 1, 1), thickness=1)
    # 叠加编号文本
    center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2
    cv2.putText(output_image, str(label), (center_x, center_y),
                cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 0), 1)  # 黑色字体

# 创建一个掩码，找到所有黑色区域（RGB: (0,0,0)）
black_mask = (output_image[:, :, 0] == 0) & (output_image[:, :, 1] == 0) & (output_image[:, :, 2] == 0)

# 将黑色区域替换为白色（RGB: (255,255,255)）
output_image[black_mask] = [255, 255, 255]

# 保存并显示结果图像
output_path = 'data/labeled_black_region_colored.png'
cv2.imwrite(output_path, output_image)

cv2.imshow('Colored Labeled Image', output_image)
cv2.waitKey(0)
cv2.destroyAllWindows()
